# -*- coding: utf-8 -*-
"""
******* 文档说明 ******
基于 tkinter 写的 GUI 图片类别标注工具
可快速查看图片、打标 并导出配置文件

调用命令：
python ImageClassify.py  图片文件夹路径  开始图片序号（可选项）

空格键：   开始、暂停播放图片
Ctrl + R： 自动播放图片顺序逆转
Ctrl + S： 保存标注标签 JSON 文件
1、2、3 ： 根据 label_mapping 打标
方向键盘： 向上、向下查看图片

Ctrl + W： 标注图片分文件夹另存，存放路径与源图片同目录

# 当前项目: ImageClassify
# 创建时间: 2019/6/30 21:25
# 开发作者: vincent
# 创建平台: PyCharm Community Edition
# 版    本: V1.0
"""
import os
import json
import time
import shutil
from tkinter import *
from PIL import Image, ImageTk
from collections import Counter


# 文件夹图片列表读取
def folder_img_list(folder_path):
    assert os.path.exists(folder_path), print('Error: 图片文件夹 【{}】不存在!!!'.format(folder_path))

    img_path_list = list()
    for root, _, img_list in os.walk(folder_path):
        for img_i in img_list:
            if img_i.endswith(('.jpg', '.JPG')):
                img_path_list.append(os.path.join(root, img_i))
    # 排序
    img_path_list = sorted(img_path_list)
    return img_path_list


# 图片查看、标注
class AppUI(object):
    def __init__(self, input_img_folder, label_mapping, scaling=0.8, sleep_time=0.5, auto_save_path=10, start_index=0):
        """
        :param input_img_folder:      图片路径
        :param label_mapping:         标签类别字典  '1': 'True', '2': 'False'}
        :param scaling:               图片显示缩放比例
        :param sleep_time:            图自动片播放时，图片间暂停时间
        :param auto_save_path:        标注结果自动保存步长
        """
        # 当前文件夹路径
        self.cwd = os.path.dirname(__file__)
        print('CurrentPath:{}'.format(self.cwd))

        # 输入标注图片文件夹路径
        self.input_img_folder = input_img_folder
        # 文件夹中所有图片路径列表
        self.img_list = folder_img_list(self.input_img_folder)
        # 图片数量
        self.img_len = len(self.img_list)
        print('ImagePath:{} num:{}'.format(self.input_img_folder, self.img_len))
        
        # 当前显示图片位置索引
        if start_index < self.img_len:
            self.img_index = start_index
        else:
            self.img_index = self.img_len - 1

        # 创建标注结果保存文件夹
        os.makedirs(os.path.join(self.cwd, 'LabelResult'), exist_ok=True)
        # 标注结果保存名称 
        json_name = 'ImageLabel---{}.json'.format(os.path.basename(self.input_img_folder))
        self.label_json_path = os.path.abspath(os.path.join(self.cwd, 'LabelResult', json_name))
        
        # 判断是否已存在此文件夹标注标签，若无则新建
        if os.path.isfile(self.label_json_path):
            self.label_dict = json.load(open(self.label_json_path, 'r', encoding='utf-8'))
            # 新增图片更新配置文件字典
            new_img_list = list(set(self.img_list) - set(self.label_dict.keys()))
            self.label_dict.update(dict(zip(new_img_list, [''] * len(new_img_list))))
        else:
            self.label_dict = dict(zip(self.img_list, ['']*self.img_len))
            
        # 自动保存步长
        self.auto_save_step = auto_save_path
        
        # 图片显示缩放比例
        self.scaling = scaling
        # 标签值
        self.label_mapping = label_mapping
        self.label_mapping_string = ' '.join(['{}: {}     '.format(key_i, self.label_mapping[key_i])
                                              for key_i in sorted(self.label_mapping)])
        # 标签值索引
        self.label_mapping_index = {value_: key_ for key_, value_ in self.label_mapping.items()}

        # 颜色列表
        color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#FF00FF', '#00FFFF',
                      '#8F0000', '#008F00', '#00008F', '#8F8F00', '#8F008F', '#008F8F']
        # 标签颜色
        label_name = [label_i[0] for label_i in sorted(self.label_mapping_index.items(), key=lambda x: x[1])]
        self.label_color = dict(zip(label_name, color_list[:len(self.label_mapping)]))

        # 自动播放
        self.stop = True
        self.forward = True
        self.sleep_time = sleep_time

        self.root = Tk()

        # root.resizable(False, False)  # 调用方法会禁止根窗体改变大小
        # 以下方法用来计算并设置窗体显示时，在屏幕中心居中
        cur_width = self.root.winfo_width()  # get current width
        cur_height = self.root.winfo_height()  # get current height
        self.scn_width, self.scn_height = self.root.maxsize()  # get screen width and height
        window_size = '+%d+%d' % ((self.scn_width - cur_width) / 2, (self.scn_height - cur_height) / 2)
        self.root.geometry(window_size)  # 设定窗口大小

        # ########################################
        # 图片显示
        self.label_img = Label(self.root, width=int(self.scn_width * self.scaling),
                               height=int(self.scn_height * self.scaling))
        self.label_img.pack(fill=X, expand=YES, side=TOP, padx=0, pady=0)

        # ########################################
        # 标签显示
        self.btn_label = Button(self.root, text="【{:^30s}】    {}".format('', self.label_mapping_string),
                                bg='white', fg='white', width=66, height=1, state=DISABLED)
        self.btn_label.bind_all('<KeyPress>', self._keyboard_cmd)

        # 图片按类别保存
        self.btn_label.bind_all('<Control-W>', self._img_re_save)
        self.btn_label.bind_all('<Control-w>', self._img_re_save)
        # 播放顺序逆转
        self.btn_label.bind_all('<Control-R>', self._auto_reverse)
        self.btn_label.bind_all('<Control-r>', self._auto_reverse)
        # 保存结果
        self.btn_label.bind_all('<Control-S>', self._label_save)
        self.btn_label.bind_all('<Control-s>', self._label_save)
        self.btn_label.pack(fill=BOTH, expand=YES, padx=0, pady=0)

        # ########################################
        # 路径显示
        self.label_path = StringVar()  # 将label标签的内容设置为字符类型，用var来接收hit_me函数的传出内容用以显示在标签上
        label_path = Entry(self.root, text=self.label_path, state='readonly', bg='green', font=('Arial', 10), width=80)
        label_path.pack(fill=X, expand=YES, side=LEFT)

        # ####################################
        # 查看下一张 (快捷键 右箭头、下箭头)
        self.btn_backward = Button(self.root, text=" 》", bg='blue', fg='white', width=1, height=2,
                                   command=self.img_show_backward)
        self.btn_backward.bind_all('<KeyPress-Right>', self.img_show_backward)
        self.btn_backward.bind_all('<KeyPress-Down>', self.img_show_backward)
        self.btn_backward.pack(fill=X, expand=YES, side=LEFT)

        # 查看上一张 (快捷键 左箭头、上箭头)
        self.btn_forward = Button(self.root, text=" 《 ", bg='blue', fg='white', width=1, height=2,
                                  command=self.img_show_forward, state=DISABLED)
        self.btn_backward.bind_all('<KeyPress-Left>', self.img_show_forward)
        self.btn_backward.bind_all('<KeyPress-Up>', self.img_show_forward)
        self.btn_forward.pack(fill=X, expand=YES, side=LEFT)

        # 创建菜单栏
        # self.create_menu()
        # self.root.title("ImageView")

        # 显示第一张图片
        self._img_show()

        self.root.mainloop()

    # 创建菜单栏
    def create_menu(self):
        menu = Menu(self.root)

        about_menu = Menu(menu, tearoff=0)
        about_menu.add_command(label="图片查看打标工具 version:1.0 开发者：陈庆伟")

        # 在菜单栏中添加以下一级菜单
        menu.add_cascade(label="关于", menu=about_menu)
        self.root['menu'] = menu
        
    # 自动查看图片方向转换
    def _auto_reverse(self, *event):
        # 输入按键
        keyboard = event[0].keysym
        del keyboard

        self.forward = not self.forward
        if self.forward:
            print('AutoRun Forward!')
        else:
            print('AutoRun Backward!')

    # 根据窗口大小计算图片 Resize大小
    def _resize(self, w, h):
        # 保持原图纵横比 Resize
        if w/h > self.scn_width/self.scn_height:
            resize_w = int(self.scn_width * self.scaling)
            resize_h = int(resize_w * h / w)
        else:
            resize_h = int(self.scn_height * self.scaling)
            resize_w = int(resize_h * w / h)
        return resize_w, resize_h

    # 显示图片
    def _img_show(self):
        img_path = self.img_list[self.img_index]
        # 显示图片路径
        self.label_path.set(img_path)
        # 打开图片
        img_open = Image.open(img_path)
        # 原始大小
        raw_size = img_open.size
        # Resize
        img_open = img_open.resize(self._resize(*raw_size))
        img = ImageTk.PhotoImage(img_open)
        self.label_img.config(image=img)
        self.label_img.image = img

        # 显示图片名称及序号
        self.root.wm_title("{}/{}  {}      {}".format(
            self.img_index+1, self.img_len, os.path.basename(img_path), raw_size))

        # 显示打标标签
        text = "【{:^30s}】    {}".format(self.label_dict[self.img_list[self.img_index]], self.label_mapping_string)
        label = self.label_dict[self.img_list[self.img_index]]
        if label == '':
            bg = 'white'
        else:
            bg = self.label_color[label]
        self.btn_label.config(text=text, bg=bg)

        # 自动保存
        if self.img_index % self.auto_save_step == 0:
            self._label_save()
            
    # 显示下张图片
    def img_show_backward(self, *event):
        if len(event) > 0:
            del event[0].keysym

        if self.img_index < self.img_len - 1:
            # 图片 索引ID +1
            self.img_index += 1
            # 显示图片
            self._img_show()

        # 无图片时按键不可用
        if self.img_index >= len(self.img_list) - 1:
            self.btn_backward.config(state=DISABLED)
        # 激活向前查看按键
        self.btn_forward.config(state=ACTIVE)

    # 显示上张图片
    def img_show_forward(self, *event):
        if len(event) > 0:
            del event[0].keysym

        if self.img_index > 0:
            # 图片 索引ID -1
            self.img_index -= 1
            # 显示图片
            self._img_show()

        # 无图片时按键不可用
        if self.img_index == 0:
            self.btn_forward.config(state=DISABLED)
        # 激活向后查看按键
        self.btn_backward.config(state=ACTIVE)

    # 键盘输入处理  打标、命令传入
    def _keyboard_cmd(self, *event):
        # 输入按键
        keyboard = event[0].keysym

        if keyboard in self.label_mapping:
            print('Label:{:10s}  {}'.format(self.label_mapping[keyboard], self.img_list[self.img_index]))
            # 加入当前图片标注信息
            self.label_dict[self.img_list[self.img_index]] = self.label_mapping[keyboard]

            # 判断是否为最后一张
            if self.img_index < self.img_len - 1:
                # 图片 索引ID +1
                self.img_index += 1
                # 显示下张图片
                self._img_show()

        if keyboard == 'space':
            # 暂停开关
            self.stop = not self.stop

            if self.forward:
                for i in range(self.img_len - self.img_index - 1):
                    # 图片 索引ID +1
                    self.img_index += 1
                    # 显示下张图片
                    self._img_show()

                    self.root.update()
                    time.sleep(self.sleep_time)

                    if self.stop:
                        break
            else:
                for i in range(self.img_index):
                    # 图片 索引ID +1
                    self.img_index -= 1
                    # 显示下张图片
                    self._img_show()

                    self.root.update()
                    time.sleep(self.sleep_time)

                    if self.stop:
                        break

    # 保存标注结果
    def _label_save(self, *event):
        if len(event) > 0:
            del event[0].keysym
            
        # 保存更新图片标签
        with open(self.label_json_path, "w", encoding='utf-8') as f_json:
            json.dump(self.label_dict, f_json, indent=4, sort_keys=True, ensure_ascii=False)
        
        print('{}LabelConfigSavePath: {}'.format(' '*20, self.label_json_path))

    # 图片标注结果另存
    def _img_re_save(self, *event):
        if len(event) > 0:
            del event[0].keysym

        # 图片分类结果保存路径
        re_save_path = '{}_Classification_{}'.format(self.input_img_folder,
                                                     time.strftime('%Y%m%d-%H%M%S', time.localtime()))
        print('Image Re Save Path:   {}'.format(re_save_path))
        print('Label:  {}'.format(Counter(self.label_dict.values())))

        # 创建各个类别文件夹
        for label_i in set(self.label_dict.values()):
            os.makedirs(os.path.join(re_save_path, label_i), exist_ok=True)

        # 各个图片按标注结果分类
        for i, img_i in enumerate(self.img_list):
            label_i = self.label_dict[img_i]
            print('\r{}/{} {:10s} {:20s}  '.format(i+1, self.img_len, label_i, img_i), end='')
            shutil.copy(img_i,
                        os.path.join(re_save_path, label_i, os.path.basename(img_i)))
        print('')

if __name__ == '__main__':
    import sys

    input_cmd = sys.argv[1:]
    print('用户输入命令：{}'.format(input_cmd))

    if len(input_cmd) > 0:
        # 第一个输入参数为 标注图片文件夹
        input_img_folder_ = input_cmd[0]

        # 第二个输入参数为 图片开始位置索引
        if len(input_cmd) > 1:
            start_index_ = max(int(input_cmd[1]) - 1, 0)
        else:
            start_index_ = 0

        # ###############################################################################################
        AppUI(
              input_img_folder=input_img_folder_,           # 图片文件路径，命令行输入   必选！
              start_index=start_index_,                     # 图片开始位置索引， 命令行输入，无输入默认为 0
              scaling=0.6,        # GUI 界面图片缩放比例
              sleep_time=0,       # 图片自动播放时间间隔
              auto_save_path=10,  # 自动保存步长
              label_mapping={'1': 'True', '2': 'False', '0': 'Unknown'}   # TODO 用户自定义标签字典
              )

    else:
        print('Error: 请输入正确的命令！！！ \nps:   '
              'python  ImageClassify.py D:\\Data\\kaggle_dog_cat\\train_')
